# Author : ZZH
# Date : 2025/2/6
"""
使用CLIP生成帧描述
"""
import torch
import clip
from PIL import Image
import os

# 加载CLIP模型
device = "cuda" if torch.cuda.is_available() else "cpu"
# model: 加载的CLIP模型，用于后续的图像和文本特征编码。
# preprocess: 预处理函数，用于将输入图像转换为模型所需的格式。
model, preprocess = clip.load("ViT-B/32", device=device)


def describe_frame(image_path):
    """
    生成单帧的文本描述
    """
    image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)

    # 预定义候选标签（可根据场景自定义）
    text_labels = [
        "一群人", "一个人在舞台上讲话",
        "一个空旷的地方", "人们在鼓掌",
        "一群学生"
    ]
    # 转换为模型需要的格式（如张量），并移动到指定设备（GPU/CPU）
    text_inputs = clip.tokenize(text_labels).to(device)

    with torch.no_grad():           # 不计算梯度
        image_features = model.encode_image(image)      # 将图像作为输入，并返回一个表示该图像特征的张量
        text_features = model.encode_text(text_inputs)    # 将文本标签作为输入，并返回一个表示这些标签特征的张量

        # 计算相似度
        # @符号表示矩阵乘法，softmax(dim=-1)：将相似度值转换为概率分布，使得所有相似度值的和为1
        similarity = (image_features @ text_features.T).softmax(dim=-1)
        # 获取相似度最高的标签索引
        best_match_idx = similarity.argmax().item()

    return text_labels[best_match_idx]


# 批量处理所有帧
frame_dir = "extracted_frames"              # 帧存储位置
descriptions = []
for frame_file in sorted(os.listdir(frame_dir)):
    if frame_file.endswith(".jpg"):         # .jpg结尾才是帧
        frame_path = os.path.join(frame_dir, frame_file)        # 路径全称
        desc = describe_frame(frame_path)
        descriptions.append(desc)
        print(f"Processed {frame_file}: {desc}")

# 保存描述结果
with open("frame_descriptions.txt", "w") as f:
    f.write("\n".join(descriptions))